import gymnasium as gym
from stable_baselines3 import PPO
from swimmer_pid_env import SwimmerPIDEnv  # 导入自定义的 Swimmer PID 环境
import numpy as np

# 主函数
def main():
    # 定义目标曲线，幅值变大，频率减慢
    def target_curve(t):
        # return np.array([20/180*3.14 * np.cos(t/0.01), 20/180*3.14 * np.sin(t/0.01)])
        return np.array([20/180*3.14, 20/180*3.14])

    env = SwimmerPIDEnv(render_mode='human')  # 使用自定义的 Swimmer PID 环境
    env.set_target_curve(target_curve)  # 设置目标曲线

    # 加载模型
    model = PPO.load("/home/sh/catkin_ws/src/ymbot_e_control/policy/ppo_swimmer_pid", device='cpu')

    # 打开文件以写入数据
    with open("/home/sh/catkin_ws/src/ymbot_e_control/data/swimmer_pid_ppo_data.txt", "w") as file:
        # 验证策略，渲染画面
        obs, _ = env.reset()
        for t in range(1000):
            # action, _states = model.predict(obs, deterministic=True)
            action = np.array([1, 0, 0, 0, 0, 0])
            obs, reward, done, truncated, info = env.step(action)
            env.render()  # 渲染画面

            # 获取当前状态和目标值
            actual_angles = env.state
            target_angles = env.setpoint
            tracking_error = target_angles - actual_angles
            kp1, ki1, kd1, kp2, ki2, kd2 = action
            control_output = env.get_control_output()

            # 输出数据到文件
            file.write(f"{t}, {actual_angles[0]}, {actual_angles[1]}, {target_angles[0]}, {target_angles[1]}, {tracking_error[0]}, {tracking_error[1]}, {kp1}, {ki1}, {kd1},  {kp2}, {ki2}, {kd2},{control_output[0]}, {control_output[1]}\n")

            if done or truncated:
                obs, _ = env.reset()  # 重新开始

    env.close()

if __name__ == '__main__':
    main()
